
def yolo_head(feats, anchors, num_classes, input_shape, calc_loss=False):
    num_anchors = len(anchors)
    # [1, 1, 1, num_anchors, 2]
    anchors_tensor = K.reshape(K.constant(anchors), [1, 1, 1, num_anchors, 2])

    # 获得x，y的网格
    # 这里拿13*13的做个例子
    # (13, 13, 1, 2)
    grid_shape = K.shape(feats)[1:3] # height, width
    grid_y = K.tile(K.reshape(K.arange(0, stop=grid_shape[0]), [-1, 1, 1, 1]),
        [1, grid_shape[1], 1, 1])
    grid_x = K.tile(K.reshape(K.arange(0, stop=grid_shape[1]), [1, -1, 1, 1]),
        [grid_shape[0], 1, 1, 1])
    grid = K.concatenate([grid_x, grid_y])
    grid = K.cast(grid, K.dtype(feats))

    # (batch_size,13,13,3,85)
    feats = K.reshape(feats, [-1, grid_shape[0], grid_shape[1], num_anchors, num_classes + 5])

    # 将预测值调成真实值
    # box_xy对应框的中心点
    # box_wh对应框的宽和高
    box_xy = (K.sigmoid(feats[..., :2]) + grid) / K.cast(grid_shape[::-1], K.dtype(feats))
    box_wh = K.exp(feats[..., 2:4]) * anchors_tensor / K.cast(input_shape[::-1], K.dtype(feats))
    box_angle = feats[...,4:5]
    box_confidence = K.sigmoid(feats[..., 5:6])
    box_class_probs = K.sigmoid(feats[..., 6:])

    # 在计算loss的时候返回如下参数
    if calc_loss == True:
        return grid, feats, box_xy, box_wh, box_angle
    return box_xy, box_wh, box_confidence, box_class_probs,box_angle


def box_iou(g, p):
    g=np.asarray(g)
    p=np.asarray(p)
    g = Polygon(g[:8].reshape((4, 2)))
    p = Polygon(p[:8].reshape((4, 2)))
    if not g.is_valid or not p.is_valid:
        return 0
    inter = Polygon(g).intersection(Polygon(p)).area
    union = g.area + p.area - inter
    iou = inter/union
    if union == 0:
        return 0
    else:
        return iou

def angle2point(b):
    # b = (cx, cy, rw, rh,angle)
    bow_x = b[0] + b[2] / 2 * math.cos(float(b[4]))
    bow_y = b[1] - b[2] / 2 * math.sin(float(b[4]))
    tail_x = b[0] - b[2] / 2 * math.cos(float(b[4]))
    tail_y = b[1] + b[2] / 2 * math.sin(float(b[4]))
    x1 = int(round(bow_x + b[3] / 2 * math.sin(float(b[4]))))
    y1 = int(round(bow_y + b[3] / 2 * math.cos(float(b[4]))))
    x2 = int(round(tail_x + b[3] / 2 * math.sin(float(b[4]))))
    y2 = int(round(tail_y + b[3] / 2 * math.cos(float(b[4]))))
    x3 = int(round(tail_x - b[3] / 2 * math.sin(float(b[4]))))
    y3 = int(round(tail_y - b[3] / 2 * math.cos(float(b[4]))))
    x4 = int(round(bow_x - b[3] / 2 * math.sin(float(b[4]))))
    y4 = int(round(bow_y - b[3] / 2 * math.cos(float(b[4]))))
    return np.array([[x1,y1],[x2,y2],[x3,y3],[x4,y4]],dtype='float32')



def yolo_loss(args, anchors, num_classes, ignore_thresh=.5, print_loss=False):

    # 一共有三层
    num_layers = 3 

    y_true = args[num_layers:]
    yolo_outputs = args[:num_layers]
    
    # 对anchor进行分层
    anchor_mask = [[6,7,8], [3,4,5], [0,1,2]] if num_layers==3 else [[3,4,5], [1,2,3]]

    # 得到input_shpae为416,416 
    input_shape = K.cast(K.shape(yolo_outputs[0])[1:3] * 32, K.dtype(y_true[0]))

    # 得到网格的shape为13,13;26,26;52,52
    grid_shapes = [K.cast(K.shape(yolo_outputs[l])[1:3], K.dtype(y_true[0])) for l in range(num_layers)]

    # 初始化loss的值
    loss = 0

    # 取出每一张图片
    # m的值就是batch_size
    # mf是bacthsize的浮点数
    m = K.shape(yolo_outputs[0])[0]
    mf = K.cast(m, K.dtype(yolo_outputs[0]))


    for l in range(num_layers):

        object_mask = y_true[l][..., 5:6]
        # 取出其对应的种类(m,13,13,3,(6+num_classes)
        true_class_probs = y_true[l][..., 6:]

        # 将yolo_outputs的特征层输出进行处理
        # grid为网格结构(13,13,1,2)，raw_pred为尚未处理的预测结果(m,13,13,3,85)
        # 还有解码后的xy，wh，(m,13,13,3,2)
        grid, raw_pred, pred_xy, pred_wh,pred_angle = yolo_head(yolo_outputs[l],
             anchors[anchor_mask[l]], num_classes, input_shape, calc_loss=True)
        
        #  (m,13,13,3,5)
        pred_box = K.concatenate([pred_xy,pred_wh,pred_angle])

        # 找到负样本群组，第一步是创建一个数组，[]
        ignore_mask = tf.TensorArray(K.dtype(y_true[0]), size=1, dynamic_size=True)
        object_mask_bool = K.cast(object_mask, 'bool')

        def loop_body(b, ignore_mask):
            # 取出第b副图内，真实存在的所有的box的参数
            # n,4
            true_box = tf.boolean_mask(y_true[l][b,...,0:4], object_mask_bool[b,...,0])
            # 计算预测结果与真实情况的iou
            # pred_box为13,13,3,5
            # 计算的结果是每个pred_box和其它所有真实框的iou
            # 13,13,3,n
            iou = box_iou(pred_box[b], true_box)
            # 13,13,3,1
            best_iou = K.max(iou, axis=-1)

            # 判断预测框的iou小于ignore_thresh则认为该预测框没有与之对应的真实框
            # 则被认为是这幅图的负样本
            ignore_mask = ignore_mask.write(b, K.cast(best_iou<ignore_thresh, K.dtype(true_box)))
            return b+1, ignore_mask
        
         # 遍历所有的图片
        _, ignore_mask = K.control_flow_ops.while_loop(lambda b,*args: b<m, loop_body, [0, ignore_mask])

        # 将每幅图的内容压缩，进行处理
        ignore_mask = ignore_mask.stack()
        ignore_mask = K.expand_dims(ignore_mask, -1)
        
        # 
        raw_true_xy = y_true[l][..., :2]*grid_shapes[l][:] - grid
        raw_true_wh = K.log(y_true[l][..., 2:4] / anchors[anchor_mask[l]] * input_shape[::-1])
        raw_angle = y_true[l][...,4:5]

        raw_true_wh = K.switch(object_mask, raw_true_wh, K.zeros_like(raw_true_wh))
        box_loss_scale = 2 - y_true[l][...,2:3]*y_true[l][...,3:4]

        xy_loss = object_mask * box_loss_scale * K.binary_crossentropy(raw_true_xy, raw_pred[...,0:2], from_logits=True)
        wh_loss = object_mask * box_loss_scale * 0.5 * K.square(raw_true_wh-raw_pred[...,2:4])
        angle_loss = object_mask * (raw_true_angle-raw_pred[...,4:5])

        confidence_loss = object_mask * K.binary_crossentropy(object_mask, raw_pred[...,5:6], from_logits=True)+ \
            (1-object_mask) * K.binary_crossentropy(object_mask, raw_pred[...,5:6], from_logits=True) * ignore_mask
        class_loss = object_mask * K.binary_crossentropy(true_class_probs, raw_pred[...,6:], from_logits=True)    

        xy_loss = K.sum(xy_loss) / mf
        wh_loss = K.sum(wh_loss) / mf
        angel_loss = K.sum(angel_loss)/mf
        confidence_loss = K.sum(confidence_loss) / mf    
        class_loss = K.sum(class_loss) / mf

        loss += xy_loss + wh_loss + angle_loss + confidence_loss + class_loss
        if print_loss:
            loss = tf.Print(loss, [loss, xy_loss, wh_loss,angel_loss, confidence_loss, class_loss, K.sum(ignore_mask)], message='loss: ')
    return loss
